# Copyright (c) 2017 Presto Labs Pte. Ltd.
# Author: xguo

import enum
import json
import logging
import time

from tornado import gen
from tornado.concurrent import Future
from tornado.ioloop import IOLoop
from tornado.websocket import websocket_connect, WebSocketClosedError

from coin.base.timestamp import get_timestamp
from coin.exchange.base.rest.task_scheduler import Scheduler
from coin.exchange.gdax.kr_rest.native_private_client import GdaxAuth

GDAX_WS_API_URL = 'wss://ws-feed.pro.coinbase.com'
GDAX_WS_SANDBOX_API_URL = 'wss://ws-feed-public.sandbox.pro.coinbase.com'


class MessageType(enum.Enum):
  ERROR = 'error'
  SUBSCRIBE = 'subscribe'
  SUBSCRIBE_RESPONSE = 'subscriptions'
  UNSUBSCRIBE = 'unsubscribe'
  HEARTBEAT = 'heartbeat'  # heartbeat
  TICKER = 'ticker'  # ticker
  SNAPSHOT = 'snapshot'  # level2
  L2UPDATE = 'l2update'  # level2
  RECEIVED = 'received'  # full
  OPEN = 'open'  # full
  DONE = 'done'  # full
  MATCH = 'match'  # full, match
  LAST_MATCH = 'last_match'  # match
  CHANGE = 'change'  # full
  ACTIVATE = 'activate'  # full, for stop order


class ChannelType(enum.Enum):
  HEARTBEAT = 'heartbeat'
  TICKER = 'ticker'
  LEVEL2 = 'level2'
  FULL = 'full'
  MATCHES = 'matches'
  USER = 'user'


class GdaxWsPrivateClient(object):
  def __init__(
      self,
      key_file,
      *,
      ioloop=None,
      notification_callback=None,
      close_callback=None,
      on_ready=None,
      timeout_sec=10,
      api_url=GDAX_WS_API_URL,
  ):

    self._ioloop = ioloop or IOLoop.current()
    self._api_url = api_url
    self._ws = None
    self._ready = False
    self._timeout = timeout_sec * 1e9
    self._gdax_auth = GdaxAuth(key_file)

    self._notification_callback = notification_callback or (lambda x: None)
    self._close_callback = close_callback or (lambda: None)
    self._on_ready = on_ready or (lambda: None)
    self._scheduler = Scheduler()

    self._last_received_message_timestamp = None
    self._pending_future = None
    self._stopped = False
    self._init_params = (ioloop, notification_callback, close_callback, timeout_sec)
    self._logger = logging.getLogger(__name__)
    self._channels = []

  @property
  def last_received_message_timestamp(self):
    return self._last_received_message_timestamp

  def is_ready(self):
    return self._ready

  def _check_future_timeout(self):
    timestamp = get_timestamp()
    fut = self._pending_future
    if fut is None:
      return

    if timestamp - fut.timestamp > self._timeout:
      fut.set_exception(TimeoutError('Websocket subscription timeout!'))

  def _check_connection(self):
    if self._last_received_message_timestamp is None:
      return

    timestamp = get_timestamp()
    diff_time = timestamp - self._last_received_message_timestamp
    if diff_time > self._timeout:
      self._logger.warning('No message received for %s seconds', diff_time / 1e9)
      self.stop()

  @gen.coroutine
  def _connect(self):
    # noinspection PyBroadException
    try:
      self._logger.info('%s try to connect to the websocket.', type(self).__name__)
      self._ws = yield websocket_connect(self._api_url)
      self._ioloop.add_callback(self._run)
      self._ready = True
      self._on_ready()
      self._logger.info('Gdax private client websocket connected.')
    except Exception:
      self._logger.exception('%s connect to websocket failed!', type(self).__name__)
      self.stop()

  @gen.coroutine
  def _run(self):
    while not self._stopped:
      msg = None
      # noinspection PyBroadException
      try:
        msg = yield self._ws.read_message()
      except Exception:
        self._logger.exception('Error reading ws message')
        self.stop()

      if msg is None:
        self._logger.error('Websocket is closed!')
        self.stop()
        return

      if len(msg) == 0:
        self._logger.info('Empty message received!')
        continue

      # noinspection PyBroadException
      try:
        self.on_message(msg)
      except Exception:
        self._logger.exception('Exception raised when process message!')

  def on_message(self, raw_msg: bytes):
    timestamp = get_timestamp()
    self._last_received_message_timestamp = timestamp
    if self._stopped:
      return

    try:
      data = json.loads(raw_msg)
    except ValueError:
      self._logger.exception('Invalid json string!, %s', raw_msg)
      return

    self._logger.debug('Receive ws msg: %s\n', data)
    msg_type = MessageType(data['type'])
    if msg_type == MessageType.SUBSCRIBE_RESPONSE:
      self._logger.info('Subscription done.\n %s', json.dumps(data, indent=2))
      self._pending_future.set_result(data)
      self._pending_future = None
    elif msg_type == MessageType.ERROR:
      self._logger.error('Error msg received %s', json.dumps(data, indent=2))
      self.stop()
    else:
      self._notification_callback(data)

  def _write_message(self, msg):
    try:
      self._ws.write_message(json.dumps(msg))
    except WebSocketClosedError as e:
      self._logger.error('Write message error, ws is close!')
      self._pending_future.set_exception(e)

  def send_request(self, msg):
    if self._ws is None:
      self._logger.error('websocket not connected yet! %s', msg)
      return
    if self._pending_future is not None:
      self._logger.error('Previous futures is not resolved yet!')
      return

    fut = Future()
    fut.timestamp = get_timestamp()
    fut.msg = msg
    self._pending_future = fut
    self._logger.debug('send request: %s\n', msg)
    self._ioloop.add_callback(self._write_message, msg)
    return fut

  def subscribe_auth_channels(self, product_ids: list, channels: list):
    timestamp = str(time.time())
    method = 'GET'
    path_url = '/users/self/verify'
    self._channels = [channel.value for channel in channels]
    msg = {
        'type': MessageType.SUBSCRIBE.value,
        'product_ids': product_ids,
        'channels': self._channels,
    }
    auth_header = self._gdax_auth.gen_auth(timestamp, method, path_url)
    auth_info = {
        'signature': auth_header['CB-ACCESS-SIGN'],
        'key': auth_header['CB-ACCESS-KEY'],
        'passphrase': auth_header['CB-ACCESS-PASSPHRASE'],
        'timestamp': auth_header['CB-ACCESS-TIMESTAMP'],
    }
    msg.update(auth_info)
    self._logger.info('Subscribe to gdax channels:\n%s', json.dumps(msg, indent=2))
    fut = self.send_request(msg)
    return fut

  def subscribe_channels(self, product_ids: list, channels: list):
    self._channels = [channel.value for channel in channels]
    msg = {
        'type': MessageType.SUBSCRIBE.value,
        'product_ids': product_ids,
        'channels': self._channels,
    }
    self._logger.info('Subscribe to gdax channels:\n%s', json.dumps(msg, indent=2))
    fut = self.send_request(msg)
    return fut

  def unsubscribe_channels(self):
    msg = {
        'type': MessageType.UNSUBSCRIBE.value,
        'channels': self._channels,
    }
    self._logger.info('Unsubscribe channels: %s', self._channels)
    fut = self.send_request(msg)
    return fut

  def start(self):
    self._ioloop.add_callback(self._connect)
    self._scheduler.add_periodic_task(1, self._check_future_timeout)
    self._scheduler.add_periodic_task(10, self._check_connection)
    self._scheduler.start()

  def stop(self):
    if self._stopped:
      return
    self._scheduler.stop()
    self._ready = False
    self._ioloop.add_callback(self._async_stop)
    self._logger.info('Gdax ws private client (%s) stopped!', id(self))

  @gen.coroutine
  def _async_stop(self):
    try:
      response = yield self.unsubscribe_channels()
      self._logger.info('All channels are unsubscribed! %s', response)
    except WebSocketClosedError:
      self._logger.error('Unsubscribe failed due to close ws. ')

    self._stopped = True
    if self._ws is not None:
      self._ws.close()
      self._ws = None
    gen.sleep(5)
    self._close_callback()

  def clone(self):
    return type(self)(*self._init_params)
